#include "ce_core.h"

#define ce_rbtree_node_init(n)                  \
    (n)->p = NULL;                              \
    (n)->l = NULL;                              \
    (n)->r = NULL;                              \
    (n)->prev = NULL;                           \
    (n)->next = NULL;                           \
    (n)->c = RED;

static ce_rbtree_node_t *
ce_rbtree_node_min(ce_rbtree_node_t *n) {
    while (n != NULL) {
        if (n->l == NULL) {
            return n;
        }

        n = n->l;
    }

    return n;
}

static ce_rbtree_node_t *
ce_rbtree_node_max(ce_rbtree_node_t *n) {
    while (n != NULL) {
        if (n->r == NULL) {
            return n;
        }

        n = n->r;
    }

    return n;
}

static void
ce_rbtree_set(ce_rbtree_t *rbt, ce_rbtree_node_t *ori, ce_rbtree_node_t *n) {
    if (n->k != ori->k) {
        return;
    }

    n->p = ori->p;
    n->l = ori->l;
    n->r = ori->r;

    if (n->p == NULL) {
        rbt->root = n;
    } else if (n->p->r == ori) {
        n->p->r = n;
    } else {
        n->p->l = n;
    }

    if (n->l != NULL) {
        n->l->p = n;
    }

    if (n->r != NULL) {
        n->r->p = n;
    }

    if (rbt->min == ori) {
        rbt->min = n;
    }

    if (rbt->max == ori) {
        rbt->max = n;
    }
}

static void
ce_rbtree_left_rotate(ce_rbtree_t *rbt, ce_rbtree_node_t *n) {
    ce_rbtree_node_t   *r, *p;

    r = n->r;
    p = n->p;

    // 节点右 child 的左子树变为节点的右子树
    n->r = r->l;
    if (n->r != NULL) {
        n->r->p = n;
    }

    // 节点的父节点变成右 child 的父节点
    r->p = p;
    if (p == NULL) { // 节点为根节点
        rbt->root = r;
    } else if (r == p->l) { // 节点为其父节点的左子树
        p->l = r;
    } else { // 节点为其父节点右子树
        p->r = r;
    }

    // 节点变为右 child 的左 child
    r->l = n;

    // 节点的父节点变为右 child
    n->p = r;
}

static void
ce_rbtree_right_rotate(ce_rbtree_t *rbt, ce_rbtree_node_t *n) {
    ce_rbtree_node_t   *l, *p;

    l = n->l;
    p = n->p;

    // 节点左 child 的右子树变为节点的左子树
    n->l = l->r;
    if (n->l != NULL) {
        n->l->p = n;
    }

    // 节点的父节点变为左 child 的父节点
    l->p = p;
    if (p == NULL) { // 节点为根节点
        rbt->root = l;
    } else if (n == p->l) { // 节点为父节点的左子树
        p->l = l;
    } else { // 节点为父节点的右子树
        p->r = l;
    }

    // 节点变为左 child 的右 child
    l->r = n;

    // 节点的父节点变为左 child
    n->p = l;
}

static void
ce_rbtree_switch_node(ce_rbtree_t *rbt, ce_rbtree_node_t n1, ce_rbtree_node_t *n2) {
    
}

static void
aw_rbtree_insert_fixup(ce_rbtree_t *rbt, ce_rbtree_node_t *n) {
    ce_rbtree_node_t       *p, *pp, *un;
    ce_rbtree_node_t       *min, *mid, *max, *tmp;
    ce_rbtree_node_t       *l, *r;

    p = n->p;
    pp = p->p;
    un = pp->l == p? pp->r: pp->l;

    while (1) {
        if (un->c == RED) { // 叔叔节点为红
            un->c = BLK;
            p->c = BLK;
            pp->c = RED;

            n = pp;
        } else { // 叔叔节点为黑
            // 对当前节点，父节点，爷节点进行排序
            min = n;
            mid = p;
            max = pp;

            // 冒泡排序
            if (min > mid) tmp = mid, mid = min, min = tmp;
            if (mid > max) tmp = max, max = mid, mid = tmp;
            if (min > mid) tmp = mid, mid = min, min = tmp;

            // 中间节点变为顶点，染为红色
            // 最小和最大节点变为左子节点和右子节点，染为黑色
            // 中间节点的左子树变为最小节点右子树，中间节点的右子树变为最大节点左子树
            l = mid->l, r = mid->r;
            mid->c = RED, mid->l = min, mid->r = max, mid->p = pp->p;
            min->c = BLK, min->r = r, min->p = mid;
            max->c = BLK, max->l = l, max->p = mid;

            if (rbt->root == pp) {
                rbt->root = n;
            }

            n = mid;
        }

        // n 为根节点
        if (rbt->root == n) {
            n->c = BLK;
            return;
        }

        // n 的父节点为黑
        if (n->c == BLK) {
            return;
        }

        // n 的父节点为红，父节点不为根节点
        p = n->p;
        pp = p->p;
        un = pp->l == p ? pp->r : pp->l;
    }
}

static void
aw_rbtree_delete_fixup(ce_rbtree_t *rbt, ce_rbtree_node_t *n) {

} 

ce_rbtree_node_t *
ce_rbtree_prev(ce_rbtree_node_t *n) {
    ce_rbtree_node_t   *prev;
    
    if (n == NULL) {
        return NULL;
    }

    // 左子树非空，返回左子树的最大子节点
    if (n->l != NULL) {
        return ce_rbtree_node_max(n->l);
    }

    // 左子树为空，返回第一个节点在其右子树的祖先
    prev = n->p;
    while (prev != NULL) {
        if (n == prev->r) {
            return prev;
        }

        n = prev;
        prev = n->p;
    }

    // 左子树为空，节点在其所有祖先节点的左子树上
    return NULL;   
}

ce_rbtree_node_t *
ce_rbtree_next(ce_rbtree_node_t *n) {
    ce_rbtree_node_t   *next;

    if (n == NULL) {
        return NULL;
    }

    // 右子树非空，返回右子树的最小子节点
    if (n->r != NULL) {
        return ce_rbtree_node_min(n->r);
    }

    // 右子树为空，返回第一个节点在其左子树的祖先
    next = n->p;
    while (next != NULL) {
        if (n == next->l) {
            return next;
        }

        n = next;
        next = n->p;
    }

    // 右子树为空，节点在其所有祖先节点的左子树上
    return NULL;
}

ce_rbtree_node_t *
ce_rbtree_get(ce_rbtree_t *rbt, ce_rbtree_key_t key) {
    ce_rbtree_node_t   *n;

    n = rbt->root;

    while (n != NULL) {
        if (n->k == key) {
            return n;
        }

        if (n->k > key) {
            n = n->l;
        } else { // n->k < key
            n = n->r;
        }
    }

    return NULL;
}

void
ce_rbtree_insert(ce_rbtree_t *rbt, ce_rbtree_node_t *n) {
    ce_rbtree_node_t       *p, *node;

    ce_rbtree_node_init(n);

    // rbtree 为空
    if (rbt->root == NULL) {
        rbt->root = n;
        rbt->max = n;
        rbt->min = n;
        n->c = BLK;
        return;
    }

    p = rbt->root;
    while (p != NULL) {
        if (p->k == n->k) { // key 相同
            ce_rbtree_set(rbt, p, n);
            if (!rbt->replace) {
                p->prev = n;
                n->next = p;
            }
            return;
        }

        node = p;

        if (p->k > n->k) {
            p = p->l;
        } else {
            p = p->r;
        }
    }

    // 在叶子节点上插入新节点
    n->p = node;
    if (n->k > node->k) {
        node->r = n;
        if (rbt->max == node) {
            rbt->max = n;
        }
    } else {
        node->l = n;
        if (rbt->min == node) {
            rbt->min = n;
        }
    }

    // 插入节点的父节点为黑节点，不需要调整
    if (n->p->c == BLK) {
        return;
    }

    aw_rbtree_insert_fixup(rbt, n);
}

void
ce_rbtree_delete(ce_rbtree_t *rbt, ce_rbtree_key_t key) {
}

void
ce_rbtree_delete_by_node(ce_rbtree_t *rbt, ce_rbtree_node_t *n) {
    ce_rbtree_node_t       *prev, *next;

    if (rbt->replace) {
        ce_rbtree_delete(rbt, n->k);
        return;
    }

    // 相同 key 不覆盖模式下，使用双向链表串联

    prev = n->prev;
    next = n->next;

    // 当前 key 下只有要删除的节点
    if (prev == NULL && next == NULL) {
        ce_rbtree_delete(rbt, n->k);
        return;
    }

    if (prev == NULL) { // 要删除节点位于双向链表头
        ce_rbtree_set(rbt, n, next);
    } else {
        prev->next = next; 
    }

    next->prev = prev;
}
